﻿#region License
// Copyright (c) 2009-2010 Topian System - http://www.topian.net
//
// Permission is hereby granted, free of charge, to any person
// obtaining a copy of this software and associated documentation
// files (the "Software"), to deal in the Software without
// restriction, including without limitation the rights to use,
// copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following
// conditions:
//
// The above copyright notice and this permission notice shall be
// included in all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
// EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
// OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
// NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
// HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
// WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
// OTHER DEALINGS IN THE SOFTWARE.
#endregion

using System;
using System.Globalization;
using System.IO;
using System.Linq;
using System.Reflection;
using System.Text;
using System.Xml.Linq;

namespace System.Diagnostics.IntelliTrace
{
    public class RegUtil
    {
        private XDocument _document;
        private Assembly _assembly;
        private CategoryAttribute[] _categories;
        private ModuleSpecificationAttribute[] _modules;
        private string _typeName;
        private MethodInfo _method;
        private string _methodName;
        private string _methodId;
        private DataQueryAttribute[] _dqa;
        private ParameterInfo[] _params;
        private bool _modified;
        private bool _register;

        private static bool TryGetCollectionPlanFile(string cultureName, out FileInfo file)
        {
            file = null;
            foreach (var suffix in new[] { "", "(x86)" })
            {
                file =
                    new FileInfo(
                        Environment.ExpandEnvironmentVariables(@"%ProgramFiles" + suffix +
                                                               @"%\Microsoft Visual Studio 10.0\Team Tools\TraceDebugger Tools\" +
                                                               cultureName + @"\CollectionPlan.xml"));
                if (file.Exists)
                    return true;
            }
            return false;
        }

        public static string VisualStudioCollectionPlanFile
        {
            get
            {
                string cultureName = CultureInfo.CurrentUICulture.TwoLetterISOLanguageName;
                FileInfo file;
                if (!TryGetCollectionPlanFile(cultureName, out file))
                {
                    FileInfo enFile;
                    if (!TryGetCollectionPlanFile("en", out enFile))
                        throw new NotSupportedException("Could not find VisualStudio 10 CollectionPlan.xml file");
                    var di = file.Directory;
                    if (!di.Exists)
                        di.Create();
                    file = enFile.CopyTo(file.FullName);
                }
                return file.FullName;
            }
        }

        #region Helpers

        private static TAttribute[] Instanciate<TAttribute>(object[] data)
            where TAttribute : Attribute
        {
            return (TAttribute[])data;
        }

        private TAttribute[] GetAttributes<TAttribute>()
            where TAttribute : Attribute
        {
            return Instanciate<TAttribute>(_assembly.GetCustomAttributes(typeof(TAttribute), true));
        }

        private static TAttribute[] GetAttributes<TAttribute>(MemberInfo member)
            where TAttribute : Attribute
        {
            return Instanciate<TAttribute>(member.GetCustomAttributes(typeof(TAttribute), true));
        }

        private static TAttribute[] GetAttributes<TAttribute>(ParameterInfo parameterInfo)
            where TAttribute : Attribute
        {
            return (TAttribute[])parameterInfo.GetCustomAttributes(typeof(TAttribute), true);
        }

        private void Update<T>(T obj, XElement collection, XName elementName
            , Func<XElement, T, bool> match
            , Action<XElement, T> init)
        {
            var element = collection.Elements(elementName).FirstOrDefault(x => match(x, obj));
            if (_register)
            {
                if (element == null)
                {
                    element = new XElement(elementName);
                    collection.Add(element);
                }
                init(element, obj);
            }
            else if (element != null)
            {
                element.Remove();
                _modified = true;
            }
        }

        private void SetAttributeValue(XElement element, string name, string value)
        {
            var oldValue = (string)element.Attribute(name);
            if (!string.Equals(value, oldValue))
            {
                element.SetAttributeValue(name, value);
                _modified = true;
            }
        }

        private void SetValue(XElement element, string value)
        {
            var oldValue = element.Value;
            if (!string.Equals(value, oldValue))
            {
                element.Value = value;
                _modified = true;
            }
        }

        private XElement EnsureChild(XElement parent, XName childName)
        {
            var child = parent.Element(childName);
            if (child == null)
            {
                child = new XElement(childName);
                parent.Add(child);
                _modified = true;
            }
            return child;
        }

        private void RemoveChild(XElement parent, XName childName)
        {
            var child = parent.Element(childName);
            if (child != null)
            {
                child.Remove();
                _modified = true;
            }
        }

        #endregion

        private void Update()
        {
            var tpp = EnsureChild(_document.Root, CollectionPlanSchema.TracePointProvider);
            var categoriesElement = EnsureChild(tpp, CollectionPlanSchema.Categories);
            var modulesElement = EnsureChild(tpp, CollectionPlanSchema.ModuleSpecifications);
            var eventsElement = EnsureChild(tpp, CollectionPlanSchema.DiagnosticEventSpecifications);

            _categories = GetAttributes<CategoryAttribute>();
            foreach (var category in _categories)
                Update(category, categoriesElement, CollectionPlanSchema.Category, IsCategory, UpdateCategory);

            _modules = GetAttributes<ModuleSpecificationAttribute>();
            foreach (var module in _modules)
                Update(module, modulesElement, CollectionPlanSchema.ModuleSpecification, IsModule, UpdateModule);

            foreach (var type in _assembly.GetTypes())
            {
                _typeName = type.FullName;
                foreach (var method in type.GetMethods())
                {
                    var @event = GetAttributes<DiagnosticEventAttribute>(method).FirstOrDefault();
                    if (@event == null)
                        continue;
                    _method = method;
                    _methodName = method.Name;
                    _params = method.GetParameters();
                    _dqa = GetAttributes<DataQueryAttribute>(method);

                    var sb = new StringBuilder();
                    sb.Append(_typeName);
                    sb.Append('.');
                    sb.Append(_methodName);
                    sb.Append('(');

                    var next = false;
                    foreach (var p in _params)
                    {
                        if (next)
                            sb.Append(',');
                        else
                            next = true;
                        sb.Append(p.ParameterType.FullName);
                    }
                    sb.Append("):");
                    sb.Append(method.ReturnType.FullName);
                    _methodId = sb.ToString();
                    Update(@event, eventsElement, CollectionPlanSchema.DiagnosticEventSpecification, IsEvent, UpdateEvent);
                }
            }
        }

        private static bool IsCategory(XElement e, CategoryAttribute c)
        {
            return c.Id == (string)e.Attribute("Id");
        }

        private void UpdateCategory(XElement e, CategoryAttribute c)
        {
            SetAttributeValue(e, "Id", c.Id);
            SetAttributeValue(e, "_locID", "category." + c.Id);
            SetValue(e, c.Description);
        }

        private static bool IsModule(XElement e, ModuleSpecificationAttribute m)
        {
            return m.Id == (string)e.Attribute("Id");
        }

        private void UpdateModule(XElement e, ModuleSpecificationAttribute m)
        {
            SetAttributeValue(e, "Id", m.Id);
            SetValue(e, Path.GetFileName(_assembly.Location));
            if (m.MinVersion != 0)
                SetAttributeValue(e, "minVersion", m.MinVersion.ToString());
            if (m.MaxVersion != 0)
                SetAttributeValue(e, "maxVersion", m.MaxVersion.ToString());
        }

        private static bool IsEvent(XElement e, DiagnosticEventAttribute a)
        {
            return e.Element(CollectionPlanSchema.SettingsName).Value == a.SettingsName;
        }


        private void UpdateEvent(XElement e, DiagnosticEventAttribute a)
        {
            SetAttributeValue(e, "enabled", a.Enabled ? "true" : "false");

            var catId = a.CategoryId;
            if (catId == null)
                switch (_categories.Length)
                {
                    case 0:
                        throw new NotSupportedException(
                            "At least a [assembly:Category] attribute should be specified");
                    case 1:
                        catId = _categories[0].Id;
                        break;
                    default:
                        throw new NotSupportedException(
                            "When no CategoryId is set on a [DiagnosticEvent] attribute, only one [assembly:Category] must be specified");
                }

            SetValue(EnsureChild(e, CollectionPlanSchema.CategoryId), catId);

            var settingsName = EnsureChild(e, CollectionPlanSchema.SettingsName);
            SetValue(settingsName, a.SettingsName);
            SetAttributeValue(settingsName, "_locID", "settingsName." + a.SettingsName);

            var settingsDescription = EnsureChild(e, CollectionPlanSchema.SettingsDescription);
            SetValue(settingsDescription, a.SettingsDescription);
            SetAttributeValue(settingsDescription, "_locID", "settingsDescription." + a.SettingsDescription);

            var bindings = EnsureChild(e, CollectionPlanSchema.Bindings);
            foreach (var binding in GetAttributes<BindingAttribute>(_method))
                Update(binding, bindings, CollectionPlanSchema.Binding, IsBinding, UpdateBinding);
        }

        private bool IsBinding(XElement e, BindingAttribute a)
        {
            return e.Element(CollectionPlanSchema.TypeName).Value == _typeName
                   && e.Element(CollectionPlanSchema.MethodName).Value == _methodName
                   && e.Element(CollectionPlanSchema.MethodId).Value == _methodId;
        }

        private void UpdateBinding(XElement dataQueriesElement, BindingAttribute bindingAttr)
        {
            var msId = bindingAttr.ModuleSpecificationId;
            if (msId == null)
            {
                switch (_modules.Length)
                {
                    case 0:
                        throw new NotSupportedException(
                            "At least a [assembly:ModuleSpecification] attribute should be specified");
                    case 1:
                        msId = _modules[0].Id;
                        break;
                    default:
                        throw new NotSupportedException(
                            "When no ModuleSpecificationId is set on a [Binding] attribute, only one [assembly:ModuleSpecification] must be specified");
                }
            }

            SetValue(EnsureChild(dataQueriesElement, CollectionPlanSchema.ModuleSpecificationId), msId);
            SetValue(EnsureChild(dataQueriesElement, CollectionPlanSchema.TypeName), _typeName);
            SetValue(EnsureChild(dataQueriesElement, CollectionPlanSchema.MethodName), _methodName);
            SetValue(EnsureChild(dataQueriesElement, CollectionPlanSchema.MethodId), _methodId);
            if (bindingAttr.ProgrammableQueryType == null)
            {
                RemoveChild(dataQueriesElement, CollectionPlanSchema.ProgrammableDataQuery);

                var shortDescription = EnsureChild(dataQueriesElement, CollectionPlanSchema.ShortDescription);
                SetValue(shortDescription, bindingAttr.ShortDescription);
                SetAttributeValue(shortDescription, "_locID", "shortDescription." + _typeName + "." + _methodName);

                var longDescription = EnsureChild(dataQueriesElement, CollectionPlanSchema.LongDescription);
                SetValue(longDescription, bindingAttr.LongDescription);
                SetAttributeValue(longDescription, "_locID", "longDescription." + _typeName + "." + _methodName);

                var dataQueries = EnsureChild(dataQueriesElement, CollectionPlanSchema.DataQueries);
                var oldCount = dataQueries.Elements().Count();
                var offset = 0;
                var count = 0;
                if (_method.ReturnType != typeof(void))
                {
                    var attr = GetAttributes<DataQueryAttribute>(_method.ReturnParameter).FirstOrDefault();
                    if (attr != null)
                    {
                        if (attr.Query == null && attr.Type == null)
                            attr.Type = _method.ReturnType;
                        UpdateDataQuery(dataQueries, ref offset, attr, -1);
                    }
                    count++;
                }
                foreach (var dqa in _dqa)
                    UpdateDataQuery(dataQueries, ref offset, dqa, 0);
                var pIndex = 0;
                foreach (var p in _params)
                {
                    var attr = GetAttributes<DataQueryAttribute>(p).FirstOrDefault();
                    if (attr != null)
                    {
                        if (string.IsNullOrEmpty(attr.Name))
                            attr.Name = p.Name;
                        if (attr.Query == null && attr.Type == null)
                            attr.Type = p.ParameterType;
                        UpdateDataQuery(dataQueries, ref offset, attr, pIndex);
                    }
                    pIndex++;
                    count++;
                }
            }
            else
            {
                RemoveChild(dataQueriesElement, CollectionPlanSchema.ShortDescription);
                RemoveChild(dataQueriesElement, CollectionPlanSchema.LongDescription);
                var pdq = EnsureChild(dataQueriesElement, CollectionPlanSchema.ProgrammableDataQuery);
                SetValue(EnsureChild(pdq, CollectionPlanSchema.ModuleName),
                         Path.GetFileName(bindingAttr.ProgrammableQueryType.Assembly.Location));
                SetValue(EnsureChild(pdq, CollectionPlanSchema.TypeName), bindingAttr.ProgrammableQueryType.FullName);
            }
        }

        private void UpdateDataQuery(XElement dataQueriesElement, ref int offset, DataQueryAttribute attr, int index)
        {
            var e = dataQueriesElement.Elements(CollectionPlanSchema.DataQuery).ElementAtOrDefault(offset);
            if (e == null)
            {
                e = new XElement(CollectionPlanSchema.DataQuery);
                dataQueriesElement.Add(e);
                _modified = true;
            }
            SetAttributeValue(e, "index", index.ToString());
            if (attr.Type == typeof(string) && attr.MaxSize == 0)
                throw new NotSupportedException("[DataQuery]: MaxSize must be specified when Type==typeof(string)");
            SetAttributeValue(e, "maxSize", attr.MaxSize.ToString());
            SetAttributeValue(e, "type", attr.Type.Name);
            SetAttributeValue(e, "name", attr.Name);
            SetAttributeValue(e, "query", attr.Query);
            SetAttributeValue(e, "_locID", _typeName + "." + _methodName + "." + attr.Name);
            offset++;
        }

        private static void Update(Assembly assembly, string collectionPlanFile, bool register)
        {
            if (collectionPlanFile == null)
                collectionPlanFile = VisualStudioCollectionPlanFile;

            var util = new RegUtil
            {
                _assembly = assembly,
                _document = XDocument.Load(collectionPlanFile),
                _register = register
            };
            util.Update();
            if (util._modified)
                util._document.Save(collectionPlanFile);
        }

        public static void Register(Assembly assembly, string collectionPlanFile = null)
        {
            Update(assembly, collectionPlanFile, true);
        }

        public static void Unregister(Assembly assembly, string collectionPlanFile = null)
        {
            Update(assembly, collectionPlanFile, false);
        }
    }
}